"""
    This script uses the Scikit-Optimize library (skopt) to optimize the buy and sell thresholds of a trading strategy
    based on the Relative Strength Index (RSI) and Moving Average Convergence Divergence (MACD) indicators. The script
    reads in a CSV file containing financial data, applies the trading strategy to generate buy and sell signals,
    and backtests the results to calculate the profit percentage. The script defines functions to populate the
    'enter_trade' and 'exit_trade' columns of the dataset, filter out data before the first buy signal and after the last
    sell signal, calculate the total profit percentage, and perform backtesting on the given dataset. Finally,
    the script uses the gp_minimize function from skopt to search for the best parameters for the trading strategy using
    a Bayesian optimization algorithm. The best parameters are used to generate the best profit and print the result.
"""
from skopt import gp_minimize
import pandas as pd
from skopt.space import Integer
from ta.trend import MACD
from ta.momentum import RSIIndicator
import numpy as np
from sklearn.linear_model import LinearRegression
from typing import List, Tuple

# Read the dataset
dataframe = pd.read_csv('datasets/SPY_DATA_20_YEARS/SPX_5min.csv')
dataframe.dropna(inplace=True)


class Backtester:
    def __init__(self, dataframe: pd.DataFrame):
        """
            Initializes a new instance of the Backtester class.
        """
        self.dataframe = dataframe.copy()
        self.generate_indicators()
        self.highest_profit = float('-inf')
        self.best_profit = 0

    search_space = [
        Integer(low=5, high=50, name='rsi_buy_threshold', dtype=int),
        Integer(low=50, high=100, name='rsi_sell_threshold', dtype=int),
    ]

    def generate_indicators(self) -> None:
        """
            Generates indicators for the dataframe.
        """
        macd_indicator = MACD(self.dataframe['close'], window_slow=26, window_fast=12, window_sign=9)
        self.dataframe['MACD'] = macd_indicator.macd()
        self.dataframe['Signal'] = macd_indicator.macd_signal()

        rsi_indicator = RSIIndicator(self.dataframe['close'], window=14)
        self.dataframe['RSI'] = rsi_indicator.rsi()

    def set_entry_conditions(self, buy_rsi: int) -> pd.DataFrame:
        """
            Sets entry conditions for trades based on the MACD and RSI indicators.

            Args:
                buy_rsi (int): The minimum RSI value required to enter a long trade.
            Returns:
                pd.DataFrame: The modified dataframe with 'enter_trade' column set to 1 where the entry conditions are met.
        """
        dataframe = self.dataframe.copy()
        macd_crosses_signal = (dataframe['MACD'].shift(1) < dataframe['Signal'].shift(1)) & (
                dataframe['MACD'] > dataframe['Signal'])
        rsi_condition = dataframe['RSI'].rolling(5).min() <= buy_rsi
        dataframe.loc[macd_crosses_signal & rsi_condition, 'enter_trade'] = 1
        return dataframe

    def set_exit_conditions(self, sell_rsi: int) -> pd.DataFrame:
        """
            Sets exit conditions for trades based on the MACD and RSI indicators.

            Args:
                sell_rsi (int): The maximum RSI value required to exit a long trade.
            Returns:
                pd.DataFrame: The modified dataframe with 'exit_trade' column set to 1 where the exit conditions are met.
        """
        dataframe = self.dataframe.copy()
        macd_crosses_under_signal = (dataframe['MACD'].shift(1) > dataframe['Signal'].shift(1)) & (
                dataframe['MACD'] < dataframe['Signal'])
        rsi_condition = dataframe['RSI'].rolling(5).max() >= sell_rsi
        dataframe.loc[macd_crosses_under_signal & rsi_condition, 'exit_trade'] = 1
        return dataframe

    @staticmethod
    def get_risk_free_rate() -> float:
        """
            Returns a constant risk-free rate.
        """
        risk_free_rate = 0.02  # Set a constant risk-free rate here
        return risk_free_rate

    def calculate_treynor_ratio(self, buy_signals: List[float], sell_signals: List[float],
                                market_returns: List[float]) -> Tuple[float, float, float]:
        """
            Calculates the treynor ratio for the given buy and sell signals and market returns.

            Args:
                buy_signals (List[float]): The list of buy signals generated by the trading strategy.
                sell_signals (List[float]): The list of sell signals generated by the trading strategy.
                market_returns (List[float]): The list of market returns corresponding to each trade.
            Returns:
                Tuple[float, float, float]: A tuple containing the negative treynor ratio, treynor ratio, and profit.
        """
        risk_free_rate = self.get_risk_free_rate()

        # Calculate excess_portfolio_returns
        excess_portfolio_returns = []
        profit = 0
        for buy, sell in zip(buy_signals, sell_signals):
            profit += sell - buy
            excess_return = (sell - buy) - risk_free_rate
            excess_portfolio_returns.append(excess_return)

        num_signal_pairs = len(buy_signals)

        # Ensure market_returns has the same length as buy_signals and sell_signals
        if len(market_returns) != num_signal_pairs:
            market_returns = market_returns[:num_signal_pairs]

        # Calculate beta
        X = np.array(market_returns).reshape(-1, 1)
        y = np.array(excess_portfolio_returns).reshape(-1, 1)
        model = LinearRegression().fit(X, y)
        beta = model.coef_[0][0]

        # Ensure you don't have a division by zero error
        if beta == 0:
            treynor_ratio = 0
        else:
            treynor_ratio = np.mean(excess_portfolio_returns) / beta

        return -treynor_ratio, profit

    @staticmethod
    def remove_extra_rows(dataframe: pd.DataFrame) -> pd.DataFrame:
        """
              Removes extra rows from the dataframe that are not part of any trade.

              Args:
                  dataframe (pandas.DataFrame): The dataframe to remove extra rows from.
              Returns:
                  pd.DataFrame: The modified dataframe with extra rows removed.
        """
        enter_trade_indices = dataframe[dataframe['enter_trade'] == 1].index
        exit_trade_indices = dataframe[dataframe['exit_trade'] == 1].index

        if len(enter_trade_indices) == 0 or len(exit_trade_indices) == 0:
            return pd.DataFrame()

        start_date = enter_trade_indices[0]
        end_date = exit_trade_indices[-1] + 1

        if end_date > start_date:
            dataframe.drop(dataframe.index[end_date:], axis=0, inplace=True)

        dataframe.drop(dataframe.index[:start_date], axis=0, inplace=True)
        dataframe.dropna(subset=['enter_trade', 'exit_trade'], inplace=True, thresh=1)

        return dataframe

    def process_trading_signals(self, dataframe: pd.DataFrame) -> Tuple[List[float], List[float]]:
        """
            Process the trading signals by identifying buy and sell signals and storing them in separate lists.

            Args:
                dataframe (pd.DataFrame): The data on which the trading signals will be processed.
            Returns:
                A tuple of two lists: the buy signals and sell signals.
        """
        buy_signals, sell_signals = [], []

        iteration = dataframe.iterrows()

        for buy_index, buy_trigger in iteration:
            if buy_trigger['enter_trade'] == 1:
                buy_signals.append(buy_trigger['close'])

                for sell_index, sell_trigger in iteration:
                    if sell_index > buy_index and sell_trigger['exit_trade'] == 1:
                        sell_signals.append(sell_trigger['close'])
                        break

        if len(buy_signals) > len(sell_signals):
            buy_signals.pop(-1)

        return buy_signals, sell_signals

    def evaluate_strategy(self, dataframe: pd.DataFrame, market_returns: list) -> tuple:
        """
            Evaluate the trading strategy by calculating the Treynor ratio and profit based on the buy and sell signals and market returns.

            Args:
                dataframe (pd.DataFrame): The data on which the trading signals will be evaluated.
                market_returns (list): The market returns over the same time period as the trading signals.
            Returns:
                A tuple of three values: the negative Treynor ratio, the positive Treynor ratio, and the profit percentage.
        """
        buy_signals, sell_signals = self.process_trading_signals(dataframe)
        neg_treynor, profit = self.calculate_treynor_ratio(buy_signals, sell_signals, market_returns)
        return neg_treynor, -neg_treynor, profit

    def find_optimal_parameters(self, params: list) -> float:
        """
            Find the optimal parameters for the trading strategy by running gp_minimize optimization algorithm.

            Args:
                params (list): A list of the RSI buy and sell thresholds to test.
            Returns:
                The negative Treynor ratio of the trading strategy.
        """
        buy_rsi, sell_rsi = params
        dataframe = self.set_entry_conditions(buy_rsi)
        dataframe = self.set_exit_conditions(sell_rsi)

        filtered_dataframe = self.remove_extra_rows(dataframe)

        if filtered_dataframe.empty:
            return 1e9

        market_returns = [0.0] * len(self.buy_signals)

        neg_treynor, treynor_ratio, profit = self.evaluate_strategy(filtered_dataframe, market_returns=market_returns)
        rounded_treynor = round(treynor_ratio, 4)
        rounded_profit = round(profit, 2)

        if treynor_ratio > self.highest_profit:
            self.highest_profit = treynor_ratio
            self.best_profit = profit
            print("Attempt with rsi_buy_threshold={}, rsi_sell_threshold={}, treynor_ratio={}, profit={}%".format(
                buy_rsi, sell_rsi, rounded_treynor, rounded_profit))
        return neg_treynor

    def optimize_trading_strategy(self) -> Tuple[list, float, str]:
        """
            Optimize the trading strategy by running the find_optimal_parameters method and returning the best parameters and evaluation metrics.

            Returns:
                A tuple of three values: the best parameters for the trading strategy, the best Treynor ratio, and a message displaying the best parameters and evaluation metrics.
        """
        result = gp_minimize(
            func=self.find_optimal_parameters,
            dimensions=self.search_space,
            n_calls=300,
            n_random_starts=20,
            random_state=42,
            n_jobs=-1
        )

        best_params = result.x
        dataframe = self.set_entry_conditions(best_params[0])
        dataframe = self.set_exit_conditions(best_params[1])

        filtered_dataframe = self.remove_extra_rows(dataframe)

        if not filtered_dataframe.empty:
            _, best_treynor, best_profit = self.evaluate_strategy(filtered_dataframe)
            message = "Best parameters: rsi_buy_threshold={}, rsi_sell_threshold={}, treynor_ratio={}, profit={}%".format(
                best_params[0], best_params[1], round(best_treynor, 4), round(best_profit, 2))
        else:
            best_treynor = self.highest_profit
            best_profit = self.best_profit
            message = "Best parameters: rsi_buy_threshold={}, rsi_sell_threshold={}, treynor_ratio={}, profit={}%".format(
                best_params[0], best_params[1], round(best_treynor, 4), round(best_profit, 2))

        return best_params, best_treynor, message

if __name__ == "__main__":
    backtester = Backtester(dataframe)

    print("Running Hyperopt. This could take a long time depending on the size of your dataframe.")
    best_params, best_profit, message = backtester.optimize_trading_strategy()

    print(message)
